#!/usr/bin/env python3
"""
BERT意图识别模型训练脚本
整合数据处理、模型训练和保存的完整流程
"""

import sys
import argparse
from pathlib import Path

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.append(str(project_root))

from src.data.data_loader import IntentDataLoader
from src.data.preprocessor import DataPreprocessor
from src.models.bert_model import create_model_from_config, ModelConfig
from src.models.trainer import TrainingPipeline
from src.utils.config import load_config, ConfigManager
from src.utils.logger import setup_logger, LogContext, PerformanceLogger
import logging


def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='BERT意图识别模型训练')
    
    parser.add_argument(
        '--config', 
        type=str, 
        default=None,
        help='配置文件路径 (YAML或JSON格式)'
    )
    
    parser.add_argument(
        '--data', 
        type=str,
        help='训练数据文件路径 (CSV格式)'
    )
    
    parser.add_argument(
        '--output-dir', 
        type=str,
        default='./models/checkpoints',
        help='模型输出目录'
    )
    
    parser.add_argument(
        '--model-name', 
        type=str,
        default='bert-base-chinese',
        help='预训练模型名称'
    )
    
    parser.add_argument(
        '--num-epochs', 
        type=int,
        default=3,
        help='训练轮数'
    )
    
    parser.add_argument(
        '--batch-size', 
        type=int,
        default=16,
        help='批次大小'
    )
    
    parser.add_argument(
        '--learning-rate', 
        type=float,
        default=2e-5,
        help='学习率'
    )
    
    parser.add_argument(
        '--max-length', 
        type=int,
        default=128,
        help='最大序列长度'
    )
    
    parser.add_argument(
        '--seed', 
        type=int,
        default=42,
        help='随机种子'
    )
    
    parser.add_argument(
        '--log-level', 
        type=str,
        default='INFO',
        choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'],
        help='日志级别'
    )
    
    parser.add_argument(
        '--no-save', 
        action='store_true',
        help='不保存训练好的模型'
    )
    
    parser.add_argument(
        '--eval-only', 
        action='store_true',
        help='只进行评估，不训练'
    )
    
    return parser.parse_args()


def load_and_preprocess_data(config, logger):
    """
    加载和预处理数据
    
    Args:
        config: 配置对象
        logger: 日志记录器
        
    Returns:
        tuple: (训练数据, 验证数据, 测试数据, 映射信息)
    """
    with LogContext(logger, 'INFO'):
        logger.info("开始数据加载和预处理")
        
        # 创建性能记录器
        perf_logger = PerformanceLogger(logger, "数据预处理")
        perf_logger.start()
        
        # 1. 加载数据
        logger.info(f"从 {config.data.data_path} 加载数据")
        data_loader = IntentDataLoader()
        df = data_loader.load_from_csv(config.data.data_path)
        
        # 验证数据
        if not data_loader.validate_data(df):
            raise ValueError("数据验证失败")
        
        # 获取数据统计
        stats = data_loader.get_data_statistics(df)
        logger.info(f"数据统计: {stats}")
        
        perf_logger.checkpoint("数据加载完成")
        
        # 2. 分割数据
        logger.info("分割数据集")
        train_df, val_df, test_df = data_loader.split_data(
            df, 
            test_size=config.data.test_size,
            val_size=config.data.val_size,
            random_state=config.data.random_state
        )
        
        perf_logger.checkpoint("数据分割完成")
        
        # 3. 数据预处理
        logger.info("开始数据预处理")
        preprocessor = DataPreprocessor(
            model_name=config.model.model_name,
            max_length=config.model.max_length
        )
        
        # 预处理训练数据
        train_encodings, train_labels, mappings = preprocessor.preprocess_data(train_df)
        
        perf_logger.checkpoint("训练数据预处理完成")
        
        # 预处理验证数据
        val_encodings = preprocessor.bert_tokenizer.tokenize_texts(val_df['text'].tolist())
        val_labels = preprocessor.label_processor.transform(val_df['intent'].tolist())
        
        perf_logger.checkpoint("验证数据预处理完成")
        
        # 预处理测试数据
        test_encodings = preprocessor.bert_tokenizer.tokenize_texts(test_df['text'].tolist())
        test_labels = preprocessor.label_processor.transform(test_df['intent'].tolist())
        
        perf_logger.checkpoint("测试数据预处理完成")
        
        # 保存映射信息
        mappings_path = Path(config.training.output_dir) / "mappings.json"
        # 确保目录存在
        mappings_path.parent.mkdir(parents=True, exist_ok=True)
        preprocessor.save_mappings(mappings, str(mappings_path))
        
        perf_logger.finish()
        
        logger.info("数据预处理完成")
        logger.info(f"训练集: {len(train_labels)}, 验证集: {len(val_labels)}, 测试集: {len(test_labels)}")
        
        return (
            (train_encodings, train_labels),
            (val_encodings, val_labels), 
            (test_encodings, test_labels),
            mappings
        )


def train_model(config, train_data, val_data, test_data, mappings, logger):
    """
    训练模型
    
    Args:
        config: 配置对象
        train_data: 训练数据
        val_data: 验证数据
        test_data: 测试数据
        mappings: 标签映射
        logger: 日志记录器
        
    Returns:
        训练结果
    """
    with LogContext(logger, 'INFO'):
        logger.info("开始模型训练")
        
        # 创建性能记录器
        perf_logger = PerformanceLogger(logger, "模型训练")
        perf_logger.start()
        
        # 1. 创建模型配置
        model_config = ModelConfig(
            model_name=config.model.model_name,
            num_labels=mappings['num_classes'],
            dropout_rate=config.model.dropout_rate,
            hidden_size=config.model.hidden_size
        )
        
        # 2. 创建模型
        logger.info(f"创建模型: {model_config.model_name}")
        model = create_model_from_config(model_config)
        
        perf_logger.checkpoint("模型创建完成")
        
        # 3. 创建tokenizer
        from transformers import AutoTokenizer
        tokenizer = AutoTokenizer.from_pretrained(config.model.model_name)
        
        # 4. 使用配置中的输出目录
        # 输出目录现在在 config.training.output_dir 中
        
        # 5. 创建训练流水线
        training_pipeline = TrainingPipeline(config.training)
        
        perf_logger.checkpoint("训练流水线创建完成")
        
        # 6. 开始训练
        logger.info("开始训练流水线")
        results = training_pipeline.run_training(
            model=model,
            tokenizer=tokenizer,
            train_data=train_data,
            eval_data=val_data,
            test_data=test_data,
            label_mappings=mappings
        )
        
        perf_logger.finish()
        
        # 7. 输出训练结果
        logger.info("训练完成")
        
        if results['training_summary']:
            summary = results['training_summary']
            logger.info(f"最佳模型: {summary.get('best_model_checkpoint', 'N/A')}")
            logger.info(f"最佳指标: {summary.get('best_metric', 'N/A')}")
            logger.info(f"最终准确率: {summary.get('final_accuracy', 'N/A')}")
        
        if results['test_result']:
            test_acc = results['test_result'].get('eval_accuracy', 0)
            logger.info(f"测试集准确率: {test_acc:.4f}")
        
        return results


def main():
    """主函数"""
    # 解析命令行参数
    args = parse_arguments()
    
    # 设置日志
    logger = setup_logger('train', log_level=args.log_level)
    
    try:
        logger.info("开始BERT意图识别模型训练")
        logger.info(f"命令行参数: {vars(args)}")
        
        # 1. 加载配置
        logger.info("加载配置")
        config_manager = ConfigManager()
        
        if args.config:
            config = config_manager.load_from_file(args.config)
        else:
            config = config_manager.load_default()
        
        # 2. 更新配置（命令行参数优先）
        updates = {}
        
        if args.data:
            updates.setdefault('data', {})['data_path'] = args.data
        
        if args.output_dir:
            updates.setdefault('training', {})['output_dir'] = args.output_dir
        
        if args.model_name:
            updates.setdefault('model', {})['model_name'] = args.model_name
        
        if args.num_epochs:
            updates.setdefault('training', {})['num_epochs'] = args.num_epochs
        
        if args.batch_size:
            updates.setdefault('training', {})['batch_size'] = args.batch_size
        
        if args.learning_rate:
            updates.setdefault('training', {})['learning_rate'] = args.learning_rate
        
        if args.max_length:
            updates.setdefault('model', {})['max_length'] = args.max_length
        
        if args.seed:
            updates.setdefault('training', {})['seed'] = args.seed
        
        if updates:
            config_manager.update_config(**updates)
            config = config_manager.get_config()
        
        logger.info("配置加载完成")
        logger.info(f"模型: {config.model.model_name}")
        logger.info(f"训练轮数: {config.training.num_epochs}")
        logger.info(f"批次大小: {config.training.batch_size}")
        logger.info(f"学习率: {config.training.learning_rate}")
        
        # 3. 加载和预处理数据
        train_data, val_data, test_data, mappings = load_and_preprocess_data(config, logger)
        
        # 4. 训练模型
        if not args.eval_only:
            results = train_model(config, train_data, val_data, test_data, mappings, logger)
            
            # 5. 保存模型 (如果需要)
            if not args.no_save:
                logger.info("训练完成，模型已保存")
            
            logger.info("🎉 训练流程全部完成！")
        else:
            logger.info("仅评估模式，跳过训练")
        
    except Exception as e:
        logger.error(f"训练过程中发生错误: {e}")
        import traceback
        logger.error(f"错误详情:\n{traceback.format_exc()}")
        sys.exit(1)


if __name__ == "__main__":
    main() 